#include <iostream>
#include <vector>
#include <thread>
#include <queue>
#include <functional>
#include <mutex>
#include <condition_variable>
#include <future>
#include <chrono>  // For measuring execution time
class ThreadPool {
public:
    ThreadPool(size_t threads) : stop(false) {
        for (size_t i = 0; i < threads; ++i) {
            workers.emplace_back([this, i] {
                thread_local std::string threadIdentifier = "Thread " + std::to_string(i + 1);

                while (true) {
                    std::function<void()> task;

                    {
                        std::unique_lock<std::mutex> lock(this->queue_mutex);
                        this->condition.wait(lock,
                            [this] { return this->stop || !this->tasks.empty(); });
                        if (this->stop && this->tasks.empty())
                            return;
                        task = std::move(this->tasks.front());
                        this->tasks.pop();
                    }

                    // Record start time
                    auto start = std::chrono::high_resolution_clock::now();

                    try {
                    //    std::cout << threadIdentifier << " is executing a task." << std::endl;
                        task();
                    } catch (const std::exception& e) {
                        std::cerr << threadIdentifier << " caught exception: " << e.what() << std::endl;
                        // Handle exceptions here
                    } catch (...) {
                        std::cerr << threadIdentifier << " caught non-standard exception." << std::endl;
                        // Handle all non-standard exceptions
                    }

                    // Record end time and calculate duration
                    auto end = std::chrono::high_resolution_clock::now();
                    auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count();

                    std::cout << threadIdentifier << " has completed a task in " << duration << " ms." << std::endl;

                    // Post-task cleanup or status update
                    // Implement necessary cleanup or updates here
                }
            });
        }
    }
#if __cplusplus >= 202002L 
template<class F, class... Args>
requires std::invocable<F, Args...>
    auto enqueue(F&& f, Args&&... args) -> std::future<decltype(std::forward<F>(f)(std::forward<Args>(args)...))>
#else
template<class F, class... Args>
    auto enqueue(F&& f, Args&&... args) -> std::future<typename std::result_of<F(Args...)>::type>
#endif
    {
#if __cplusplus >= 202002L 
        // Since C++20, we can use the `std::invocable` concept to check if `f` is callable with `args`
        using return_type = decltype(std::forward<F>(f)(std::forward<Args>(args)...));
        auto task = std::make_shared<std::packaged_task<return_type()>>(
                // Forwarding the function and arguments to the lambda
                [f = std::forward<F>(f), ...args = std::forward<Args>(args)]() mutable {
                    // Invoke the function with the arguments
                    return f(args...);
                }
            );
#else
        using return_type = typename std::result_of<F(Args...)>::type;
        auto task = std::make_shared< std::packaged_task<return_type()> >(
                std::bind(std::forward<F>(f), std::forward<Args>(args)...)
            );
#endif
        auto res = task->get_future();
            {
                std::unique_lock<std::mutex> lock(queue_mutex);

                // Don't allow enqueueing after stopping the pool
                if(stop.load())
                    throw std::runtime_error("enqueue on stopped ThreadPool");

                tasks.emplace([task]() { (*task)(); });
            }
            condition.notify_one();
            return res;
    }
    void shutdown() {
        stop = true;
        condition.notify_all();  // Wake up all threads to check `stop`
    }
    ~ThreadPool() {
        stop = true;
        condition.notify_all();  // Wake up all waiting threads.

        for (auto& worker : workers) {
            if (worker.joinable())
                worker.join();  // Wait for the thread to finish
        }
    }



private:
    // Need to keep track of threads so we can join them
    std::vector<std::thread> workers;
    // The task queue
    std::queue<std::function<void()>> tasks;

    // Synchronization
    std::mutex queue_mutex;
    std::condition_variable condition;
    std::atomic<bool> stop;
};

// Example usage
int main() {
    ThreadPool pool(8);  // Create a ThreadPool with 4 worker threads
    std::vector <std::future<int>> results; 
    for (int i = 0; i < 8; ++i) {
        results.emplace_back(
            pool.enqueue([i] {
                std::this_thread::sleep_for(std::chrono::seconds(1));
                return i * i;
            })
        );
    }
    // Print results
    for (size_t i = 0; i < results.size(); ++i) {
        std::cout << "Result of task " << i << " is " << results[i].get() << std::endl;
    }


    std::cout << "All tasks completed." << std::endl;   
    pool.shutdown();  // Ensure all threads are stopped before exiting
    return 0;
}